package cgroups

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	cgroupRoot = "/sys/fs/cgroup"
)

type CgroupManager struct {
	cgroupPath string
}

func NewCgroupManager(cgroupName string) *CgroupManager {
	return &CgroupManager{
		cgroupPath: filepath.Join(cgroupRoot, cgroupName),
	}
}

// Create方法创建cgroup目录
func (cg *CgroupManager) Create() error {
	if err := os.Mkdir(cg.cgroupPath, 0755); err != nil && !os.IsExist(err) {
		return fmt.Errorf("failed to create cgroup directory: %w", err)
	}
	return nil
}

/*
SetLimits方法设置cgroup的资源限制

	cpuMax: 带宽限制 (Bandwidth), 限制该 cgroup 在任何 100ms (100000us) 周期内最多使用 20ms (20000us) 的 CPU 时间 (即 20% CPU 限制)(cpu.max)
	memoryMax: # 限制该 cgroup 的内存使用量(memory.max)
	swapMax: 限制该 cgroup 的交换空间使用量(memory.swap.max)
	pidsMax: 限制该 cgroup 的进程数(pids.max)
	cpuWeight:权重更高的值意味着在 CPU 竞争时获得更多份额(cpu.weight)
	Controller: 控制组类型(cgroup.subtree_control)
		# 启用内存和 pids 控制器供当前 cgroup 的子 cgroup 使用 (前面加 '+'号)
*/
func (cg *CgroupManager) SetLimits(cpuMax string, memoryMaxMB int, pidsMax int) error {
	// 设置 CPU 限制 (格式: <max> <period>)
	cpuMaxValue := cpuMax + " 100000"
	if err := writeFile(filepath.Join(cg.cgroupPath, "cpu.max"), cpuMaxValue); err != nil {
		return fmt.Errorf("failed to set CPU limit: %w", err)
	}

	// 设置内存限制
	memoryBytes := memoryMaxMB * 1024 * 1024
	if err := writeFile(filepath.Join(cg.cgroupPath, "memory.max"),
		strconv.Itoa(memoryBytes)); err != nil {
		return fmt.Errorf("failed to set memory limit: %w", err)
	}

	if err := writeFile(filepath.Join(cg.cgroupPath, "memory.swap.max"), "0"); err != nil {
		return fmt.Errorf("failed to set swap limit: %w", err)
	}

	// 设置进程数限制
	if err := writeFile(filepath.Join(cg.cgroupPath, "pids.max"),
		strconv.Itoa(pidsMax)); err != nil {
		return fmt.Errorf("failed to set PID limit: %w", err)
	}

	return nil
}

func (cg *CgroupManager) AddProcess(pid int) error {
	return writeFile(filepath.Join(cg.cgroupPath, "cgroup.procs"),
		strconv.Itoa(pid))
}

// Destroy方法确保可靠清理
func (cg *CgroupManager) Destroy() error {
	procsFile := filepath.Join(cg.cgroupPath, "cgroup.procs")
	rootProcsFile := filepath.Join(cgroupRoot, "cgroup.procs")

	// 读取当前cgroup中的所有进程PID
	data, err := os.ReadFile(procsFile)
	if err != nil {
		// 文件不存在或读取失败时直接删除cgroup
		return os.RemoveAll(cg.cgroupPath)
	}

	pids := strings.Split(strings.TrimSpace(string(data)), "\n")
	for _, pidStr := range pids {
		if pidStr == "" {
			continue
		}
		// 将进程迁移到根cgroup
		if err := writeFile(rootProcsFile, pidStr); err != nil {
			// 进程可能已退出，忽略NOT_FOUND错误
			if !os.IsNotExist(err) {
				return fmt.Errorf("failed to move process %s: %w", pidStr, err)
			}
		}
	}

	// 删除cgroup目录
	return os.RemoveAll(cg.cgroupPath)
}

/*cgroups目录结构下统计文件
* memory.current: 当前内存使用量。
* memory.stat: 详细内存统计（缓存、RSS 等）。
* cpu.stat: CPU 使用统计（用户态、内核态时间，周期数等）。
* pids.current: 当前进程数。
 */
// 获取CPU使用时间（纳秒）
func (cg *CgroupManager) GetCPUUsage() (uint64, error) {
	data, err := os.ReadFile(filepath.Join(cg.cgroupPath, "cpu.stat"))
	if err != nil {
		return 0, err
	}

	for _, line := range strings.Split(string(data), "\n") {
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}
		if fields[0] == "usage_usec" {
			usec, err := strconv.ParseUint(fields[1], 10, 64)
			if err != nil {
				return 0, fmt.Errorf("parse usage_usec failed: %w", err)
			}
			// 微秒 → 纳秒 (1μs = 1000ns)
			return usec * 1000, nil
		}
	}
	return 0, fmt.Errorf("usage_usec not found in cpu.stat")
}

// 获取内存使用量（字节）
func (cg *CgroupManager) GetMemoryUsage() (uint64, error) {
	data, err := os.ReadFile(filepath.Join(cg.cgroupPath, "memory.current"))
	if err != nil {
		return 0, err
	}
	return strconv.ParseUint(strings.TrimSpace(string(data)), 10, 64)
}

func writeFile(path, value string) error {
	return os.WriteFile(path, []byte(value), 0644)
}
